/*
 * Forward the HTTPS request from libcurl to instance vm.
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */
#include <stdio.h>
#include <stdint.h>
#include <errno.h>
#include <sys/un.h>
#include <sys/socket.h>
#include <linux/vm_sockets.h>
#include <unistd.h>
#include <signal.h>
#include "enclave_proxy.h"

#define BUF_SIZE                0x1000
#define MAX_BACKLOG             1

volatile sig_atomic_t loop_flag;

int un_server_init(struct connect_info *conn, const char *filepath)
{
    struct sockaddr_un server_sockaddr = {0};
    int un_sockaddr_len;
    int rc;

    if (sizeof(server_sockaddr.sun_path) <= strlen(filepath))
        return PX_INVALID_ARGUMENT;

    /* Create a unix domain stream socket */
    conn->unix_info.server_sockfd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (conn->unix_info.server_sockfd < 0)
        return PX_SOCKET_ERROR;

    /* Set up the UNIX sockaddr structure and give it a filepath to bind to. */
    server_sockaddr.sun_family = AF_UNIX;
    strncpy(server_sockaddr.sun_path, filepath, strlen(filepath));
    un_sockaddr_len = sizeof(server_sockaddr);

    unlink(filepath);

    rc = bind(conn->unix_info.server_sockfd, (struct sockaddr *)&server_sockaddr,
        un_sockaddr_len);
    if (rc < 0) {
        close(conn->unix_info.server_sockfd);
        return PX_BIND_ERROR;
    }

    /* Listen for any client sockets */
    rc = listen(conn->unix_info.server_sockfd, MAX_BACKLOG);
    if (rc < 0) {
        close(conn->unix_info.server_sockfd);
        return PX_LISTEN_ERROR;
    }
    fprintf(stdout, "unix socket listening...\n");

    return PX_NO_ERROR;
}

int vsock_client_init(connect_info *conn, int vsock_cid, int vsock_port)
{
    struct sockaddr_vm client_vsockaddr = {0};
    int rc;
    int timeout = VSOCK_TIMEOUT;

    conn->vsock_info.client_sockfd = socket(AF_VSOCK, SOCK_STREAM, 0);
    if (conn->vsock_info.client_sockfd < 0)
        return PX_SOCKET_ERROR;

    client_vsockaddr.svm_family = AF_VSOCK;
    client_vsockaddr.svm_port = vsock_port;
    client_vsockaddr.svm_cid = vsock_cid;

    do {
        /* Connect to the vsock channel in the parent vm.(wait for 5s to timeout) */
        rc = connect(conn->vsock_info.client_sockfd,
            (struct sockaddr *)&client_vsockaddr, sizeof(struct sockaddr_vm));
        if (rc != 0)
            sleep(1);
        else
            break;
        timeout--;
        fprintf(stderr, "timeout count down is %ds\n", timeout);
    } while (timeout > 0);
    if (rc != 0) {
        close(conn->vsock_info.client_sockfd);
        return PX_CONNECT_ERROR;
    }

    return PX_NO_ERROR;
}

static int recv_data(int sockfd, uint8_t *buf, ssize_t len, int flags,
                     ssize_t *bytes)
{
    ssize_t cur_bytes;
    int count;
    fd_set readfs;
    struct timeval timeout;
    int rc = PX_NO_ERROR;
    int retry_time = 0;

    do {
        /* Make select wait up to 100ms for incoming data. */
        FD_ZERO(&readfs);
        FD_SET(sockfd, &readfs);
        timeout.tv_sec = 0;
        timeout.tv_usec = RECV_DATA_TIMEOUT_US;
        count = select(sockfd + 1, &readfs, NULL, NULL, &timeout);
        if (count < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
            retry_time++;
            continue;
        } else if (count < 0) {
            rc = PX_INTER_ERROR;
            break;
        } else if (count == 0) {
            /* Timeout */
            break;
        }

        cur_bytes = recv(sockfd, buf + *bytes, len - *bytes, flags);
        /* retry */
        if (cur_bytes < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
            retry_time++;
            continue;
        } else if (cur_bytes < 0) {
            /* ERROR */
            rc = PX_INTER_ERROR;
            break;
        } else if (cur_bytes == 0) {
            /* NO data left */
            rc = PX_NO_DATA;
            break;
        }

        *bytes += cur_bytes;
        /* Buf full */
        if (*bytes == len) {
            break;
        } else if (*bytes > len) {
            rc = PX_INTER_ERROR;
            break;
        }

        retry_time = 0;
    } while (count > 0 || retry_time < RETRY_MAX_TIMES);

    return rc;
}

static int transfer_data(int srcfd, int dstfd)
{
    uint8_t buf[BUF_SIZE];
    ssize_t recv_bytes;
    ssize_t send_bytes;
    int rc;

    recv_bytes = 0;
    rc = recv_data(srcfd, buf, sizeof(buf), 0, &recv_bytes);
    if ((recv_bytes == 0 && rc == PX_NO_DATA) ||
        (rc != PX_NO_ERROR && rc != PX_NO_DATA))
            return rc;

    if (recv_bytes != 0) {
        send_bytes = send(dstfd, buf, recv_bytes, MSG_NOSIGNAL);
        if (send_bytes != recv_bytes)
            return PX_INTER_ERROR;
    }

    return PX_NO_ERROR;
}

static int handle_routine(const struct connect_info *conn)
{
    int rc;

    while (loop_flag) {
        /* unix socket to vsock */
        rc = transfer_data(conn->unix_info.client_sockfd, conn->vsock_info.client_sockfd);
        if (rc == PX_NO_DATA)
            return PX_NO_CONNECT;
        else if (rc != PX_NO_ERROR)
            return rc;

        /* vsock to unix socket */
        rc = transfer_data(conn->vsock_info.client_sockfd, conn->unix_info.client_sockfd);
        if (rc != PX_NO_ERROR && rc != PX_NO_DATA)
            return rc;
    }

    return PX_NO_ERROR;
}

static void *process_request(void *arg)
{
    struct sockaddr_un client_sockaddr = {0};
    int un_sockaddr_len;
    fd_set readfs;
    struct timeval timeout;
    struct connect_info *conn = (struct connect_info *)arg;
    int rc;

    while (loop_flag) {
        timeout.tv_sec = 0;
        timeout.tv_usec = CONNECT_TIMEOUT_US;

        FD_ZERO(&readfs);
        FD_SET(conn->unix_info.server_sockfd, &readfs);

        rc = select(conn->unix_info.server_sockfd + 1, &readfs, NULL, NULL, &timeout);
        if (rc <= 0) {
            if (rc == 0 || errno == EAGAIN || errno == EWOULDBLOCK)
                continue;
            else
                break;
        } else if (FD_ISSET(conn->unix_info.server_sockfd, &readfs)) {
            /* Accept an incoming unix socket connection */
            un_sockaddr_len = sizeof(client_sockaddr);
            memset(&client_sockaddr, 0, un_sockaddr_len);
            conn->unix_info.client_sockfd = accept(conn->unix_info.server_sockfd,
                (struct sockaddr *)&client_sockaddr, &un_sockaddr_len);
            if (conn->unix_info.client_sockfd < 0) {
                fprintf(stderr, "unix server accept error!\n");
                break;
            }

            /* Get the name of the connected unix socket */
            rc = getpeername(conn->unix_info.client_sockfd,
                (struct sockaddr *)&client_sockaddr, &un_sockaddr_len);
            if (rc < 0) {
                fprintf(stderr, "getpeername error!\n");
                break;
            } else {
                fprintf(stderr, "Client socket filepath: %s!\n",
                    client_sockaddr.sun_path);
            }

            rc = handle_routine(conn);
            if (rc != PX_NO_ERROR && rc != PX_NO_CONNECT) {
                fprintf(stderr, "handle_routine error: %d\n", rc);
                break;
            }
        }
    }
    close(conn->unix_info.server_sockfd);
    close(conn->vsock_info.client_sockfd);
    memset(conn, 0, sizeof(struct connect_info));
    unlink(CURL_SOCKET_PATH);
    return NULL;
}

int create_request_process_thread(struct connect_info *conn)
{
    int rc;

    rc = pthread_create(&(conn->thread_id), NULL, process_request, conn);
    if (rc != 0) {
        close(conn->unix_info.server_sockfd);
        close(conn->vsock_info.client_sockfd);
        return PX_PTHREAD_ERROR;
    }

    return PX_NO_ERROR;
}

int setup_proxy(struct connect_info *conn, int vsock_cid, int vsock_port)
{
    int rc;

    if ((conn == NULL) || (vsock_cid <= VMADDR_CID_HOST))
        return PX_INVALID_ARGUMENT;

    /* unix socket */
    rc = un_server_init(conn, CURL_SOCKET_PATH);
    if (rc != PX_NO_ERROR)
        return rc;

    /* vsock */
    rc = vsock_client_init(conn, vsock_cid, vsock_port);
    if (rc != PX_NO_ERROR) {
        close(conn->unix_info.server_sockfd);
        return rc;
    }

    loop_flag = 1;
    rc = create_request_process_thread(conn);
    if (rc != PX_NO_ERROR)
        return rc;

    return PX_NO_ERROR;
}

void close_proxy(struct connect_info *conn)
{
    if (conn == NULL)
        return;

    loop_flag = 0;
    pthread_join(conn->thread_id, NULL);
    close(conn->unix_info.server_sockfd);
    close(conn->vsock_info.client_sockfd);
    memset(conn, 0, sizeof(struct connect_info));

    /* remove unix socket file */
    unlink(CURL_SOCKET_PATH);
}
